package com.xiao.xin.core.service.Impl;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson2.JSON;
import com.xiao.xin.core.common.Constants;
import com.xiao.xin.core.domain.ComfyuiTask;
import com.xiao.xin.core.service.RedisService;
import lombok.RequiredArgsConstructor;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Service;

import java.time.Duration;

import static com.xiao.xin.core.common.Constants.DISTRIBUTED_QUEUE_KEY;
import static com.xiao.xin.core.common.Constants.TASK_KEY_PREFIX;

@RequiredArgsConstructor
@Service
public class RedisServiceImpl implements RedisService {

    private final StringRedisTemplate redisTemplate;
    private final static String DISTRIBUTED_ID_KEY = "DISTRIBUTED_ID";

    @Override
    public ComfyuiTask addQueueTask(ComfyuiTask comfyuiTask) {
        Long idx = redisTemplate.opsForValue().increment(DISTRIBUTED_ID_KEY);
        comfyuiTask.setIndex(idx);

        redisTemplate.opsForZSet().add(Constants.DISTRIBUTED_QUEUE_KEY, comfyuiTask.getId(), idx);

        redisTemplate.opsForValue().set(Constants.TASK_KEY_PREFIX + comfyuiTask.getId(), JSON.toJSONString(comfyuiTask));
        return comfyuiTask;
    }

    @Override
    public ComfyuiTask popQueueTask() {
        Long size = redisTemplate.opsForZSet().size(Constants.DISTRIBUTED_QUEUE_KEY);
        if (size > 0) {
            ZSetOperations.TypedTuple<String> task = redisTemplate.opsForZSet().popMin(Constants.DISTRIBUTED_QUEUE_KEY);
            if (task != null) {
                String taskId = task.getValue();
                String json = redisTemplate.opsForValue().get(Constants.TASK_KEY_PREFIX + taskId);
                return JSON.parseObject(json, ComfyuiTask.class);
            }
        }

        return null;
    }

    @Override
    public void addStartedTask(String promptId, ComfyuiTask task) {
        redisTemplate.opsForValue().set(Constants.RUN_TASK_KEY + promptId, JSON.toJSONString(task), Duration.ofMinutes(10));
    }

    @Override
    public ComfyuiTask getStartedTask(String promptId) {
        String json = redisTemplate.opsForValue().get(Constants.RUN_TASK_KEY + promptId);
        if (StrUtil.isNotEmpty(json)) {
            return JSON.parseObject(json, ComfyuiTask.class);
        }
        return null;
    }

    @Override
    public boolean hasQueueTask() {
        return redisTemplate.opsForZSet().size(DISTRIBUTED_QUEUE_KEY) > 0;
    }
}
